###############################################################################   
#### Replication Materials                                                 #### 
#### Kim, Nakka, Gopal, Desmrais, Mancinelli, Harden, Ko, Boehmke. 2021.   ####
#### Attention to the COVID-19 pandemic on Twitter:                        ####
#### Partisan differences among U.S. state legislators                     ####
#### Legislative Studies Quarterly                                         ####
###############################################################################  


###############################################################################
################################### Set Up ####################################
###############################################################################

# packages -------------------------

lapply(c('readr', 'stargazer', 'texreg', 'xtable', 'lme4', 'glmmTMB','plm',
         'lmtest', 'interplot','dummies','effects', 'ggthemes','caret', 
         'ggeffects', 'scales','ggrepel','sandwich','reshape2'), 
       require, 
       character.only = TRUE)

# read data set for regression -------------------------

path_data <- '/Users/taegyoon/Google Drive/spap_state/spap_state_attention/data/'

df <- read_csv(paste0(path_data, "spap_state_attention_regression.csv"))
df <- subset(df, select = -c(X1) )
df_uni <- df[which(df$unified_gov == TRUE), ]
df_spt <- df[which(df$unified_gov == FALSE), ]


# set the panel structure -------------------------

plm_df_uni <- pdata.frame(x = df_uni, 
                          index = c('user.screen_name', 'week'))
plm_df_spt <- pdata.frame(x = df_spt, 
                          index = c('user.screen_name', 'week'))


###############################################################################
############################ Fit Regression Model #############################
###############################################################################

uni_all_inter <- plm(covid_relevant_1_log ~ 
                       party_new*state_case_pop + 
                       party_new*state_death_pop +
                       party_new*national_case_pop + 
                       party_new*national_death_pop +
                       party_new*state_covid_policy + 
                       week_new +
                       week_new_2 +
                       week_new_3,            
                     data = plm_df_uni,
                     index = c('user.screen_name', 'week'),
                     model = 'random', 
                     effect = 'twoways')


###############################################################################
###################### Generate Effect Plots: Figure S11 ######################
###############################################################################

# variance-covariance matrix for adjusted SE -------------------------

vcov <- vcovHC(uni_all_inter, 
               method = 'arellano')

# state case -------------------------

state_case_values <- seq(round(min(df_uni$state_case_pop), 1),
                         round(max(df_uni$state_case_pop), 1),
                         0.1)
state_case_pop_effect <- Effect(focal.predictors = c('state_case_pop',
                                                     'party_new'), 
                                mod = uni_all_inter,
                                xlevels = list(state_case_pop = state_case_values),
                                given.values = c(
                                  state_death_pop = median(df_uni$state_death_pop),
                                  national_case_pop = median(df_uni$national_case_pop),
                                  national_death_pop = median(df_uni$national_death_pop),
                                  state_covid_policy = median(df_uni$state_covid_policy),
                                  week_new = median(df_uni$week_new),
                                  week_new_2 = median(df_uni$week_new_2),
                                  week_new_3 = median(df_uni$week_new_3)
                                  )
                                )

state_case_me_se_dem <- sqrt(vcov[3, 3]) # SE for marginal effect of predictor
state_case_me_se_rep <- sqrt(vcov[3, 3] 
                             + (1)^2 * vcov[11, 11] 
                             + 2 * 1 * vcov[3, 11]) # SE for marginal effect of predictor

state_case_me_ci_dem <- cbind(state_case_values * (1.96 * state_case_me_se_dem), 
                              state_case_values * -(1.96 * state_case_me_se_dem))
state_case_me_ci_rep <- cbind(state_case_values * (1.96 * state_case_me_se_rep), 
                              state_case_values * -(1.96 * state_case_me_se_rep))

state_case_preds <- state_case_pop_effect$fit # in the order of Dem, Neither, Rep
state_case_upper <- state_case_pop_effect$fit + c(state_case_me_ci_dem[, 1], 
                                                  state_case_me_ci_rep[, 1])
state_case_lower <- state_case_pop_effect$fit + c(state_case_me_ci_dem[, 2], 
                                                  state_case_me_ci_rep[, 2])
state_case_final <- data.frame(Party = c(rep('Democrat', length(state_case_values)), 
                                         rep('Republican', length(state_case_values))), 
                               X = as.numeric(state_case_values), 
                               Y = as.numeric(state_case_preds), 
                               U = as.numeric(state_case_upper), 
                               L = as.numeric(state_case_lower),
                               Y_linear = exp(as.numeric(state_case_preds)) - 1, 
                               U_linear = exp(as.numeric(state_case_upper)) - 1, 
                               L_linear = exp(as.numeric(state_case_lower)) - 1)

state_case_final$Party <- as.factor(state_case_final$Party)
state_case_final$Party <- factor(state_case_final$Party, 
                                 levels = rev(levels(state_case_final$Party)))

ggplot(state_case_final, 
       aes(x = X, 
           y = Y_linear, 
           color = Party))+ 
  theme_calc() +
  geom_line() +
  geom_ribbon(aes(ymin = L_linear, ymax = U_linear), 
              linetype = 2, 
              alpha = 0.0) +
  theme(text = element_text(size = 12.5), 
        plot.title = element_text(hjust = 0.5),
        plot.margin = grid::unit(c(5, 5, 5, 5), "mm"),
        legend.title = element_blank(),
        legend.position = 'bottom',
        legend.direction = 'horizontal') +
  scale_color_manual(values = c('darkred', 'dodgerblue4')) +
  xlab('\nWeekly Count of State New Cases (per 10k)') + 
  ylab('Count of COVID-19 Tweets\n') +
  geom_rug(data = df_uni, 
           aes(x = state_case_pop), 
           inherit.aes = FALSE, 
           color = 'gray60') +
  scale_shape_cleveland(overlap = FALSE)

ggsave(paste0(rep_path, 'FigS11a.png'),
       dpi = 600,
       width = 5,
       height = 3.5,
       units = 'in')

# state death -------------------------

state_death_values <- seq(round(min(df_uni$state_death_pop), 1),
                          round(max(df_uni$state_death_pop), 1),
                          0.1)
state_death_pop_effect <- Effect(focal.predictors = c('state_death_pop',
                                                      'party_new'), 
                                 mod = uni_all_inter,
                                 xlevels = list(state_death_pop = state_death_values),
                                 given.values = c(
                                   state_case_pop = median(df_uni$state_case_pop),
                                   national_case_pop = median(df_uni$national_case_pop),
                                   national_death_pop = median(df_uni$national_death_pop),
                                   state_covid_policy = median(df_uni$state_covid_policy),
                                   week_new = median(df_uni$week_new),
                                   week_new_2 = median(df_uni$week_new_2),
                                   week_new_3 = median(df_uni$week_new_3)
                                   )
                                 )

state_death_me_se_dem <- sqrt(vcov[4, 4]) # SE for marginal effect of predictor
state_death_me_se_rep <- sqrt(vcov[4, 4] + 
                                (1)^2 * vcov[12, 12] + 
                                2 * 1 * vcov[4, 12]) # SE for marginal effect of predictor

state_death_me_ci_dem <- cbind(state_death_values * (1.96 * state_death_me_se_dem), 
                               state_death_values * -(1.96 * state_death_me_se_dem))
state_death_me_ci_rep <- cbind(state_death_values * (1.96 * state_death_me_se_rep), 
                               state_death_values * -(1.96 * state_death_me_se_rep))

state_death_preds <- state_death_pop_effect$fit # in the order of Dem, Neither, Rep
state_death_upper <- state_death_pop_effect$fit + c(state_death_me_ci_dem[, 1], 
                                                    state_death_me_ci_rep[, 1])
state_death_lower <- state_death_pop_effect$fit + c(state_death_me_ci_dem[, 2], 
                                                    state_death_me_ci_rep[, 2])
state_death_final <- data.frame(Party = c(rep('Democrat', length(state_death_values)), 
                                          rep('Republican',length(state_death_values))), 
                                X = as.numeric(state_death_values), 
                                Y = as.numeric(state_death_preds), 
                                U = as.numeric(state_death_upper), 
                                L = as.numeric(state_death_lower),
                                Y_linear = exp(as.numeric(state_death_preds)) - 1, 
                                U_linear = exp(as.numeric(state_death_upper)) - 1, 
                                L_linear = exp(as.numeric(state_death_lower)) - 1)

state_death_final$Party <- as.factor(state_death_final$Party)
state_death_final$Party <- factor(state_death_final$Party, 
                                  levels = rev(levels(state_death_final$Party)))

ggplot(state_death_final, 
       aes(x = X, 
           y = Y_linear, 
           color = Party)) + 
  theme_calc() +
  geom_line() +
  geom_ribbon(aes(ymin = L_linear, 
                  ymax = U_linear), 
              linetype = 2, 
              alpha = 0.0) +
  theme(text = element_text(size = 12.5), 
        plot.title = element_text(hjust = 0.5),
        plot.margin = grid::unit(c(5, 5, 5, 5), "mm"),
        legend.title = element_blank(),
        legend.position = 'bottom',
        legend.direction = 'horizontal') +
  scale_color_manual(values = c('darkred', 'dodgerblue4')) +
  xlab('\nWeekly Count of State New Deaths (per 10k)') + 
  ylab('Count of COVID-19 Tweets\n') +
  geom_rug(data = df_uni, 
           aes(x = state_death_pop), 
           inherit.aes = FALSE, 
           color = 'gray60') +
  scale_shape_cleveland(overlap = FALSE)

ggsave(paste0(rep_path, 'FigS11b.png'),
       dpi = 600,
       width = 5,
       height = 3.5,
       units = 'in')


# national case -------------------------

national_case_values <- seq(round(min(df_uni$national_case_pop), 1),
                            round(max(df_uni$national_case_pop), 1),
                            0.1)
national_case_pop_effect <- Effect(focal.predictors = c('national_case_pop',
                                                        'party_new'), 
                                   mod = uni_all_inter,
                                   xlevels = list(national_case_pop = national_case_values),
                                   given.values = c(
                                     state_case_pop = median(df_uni$state_case_pop),
                                     state_death_pop = median(df_uni$state_death_pop),
                                     national_death_pop = median(df_uni$national_death_pop),
                                     state_covid_policy = median(df_uni$state_covid_policy),
                                     week_new = median(df_uni$week_new),
                                     week_new_2 = median(df_uni$week_new_2),
                                     week_new_3 = median(df_uni$week_new_3)
                                     )
                                   )

national_case_me_se_dem <- sqrt(vcov[5, 5]) # SE for marginal effect of predictor
national_case_me_se_rep <- sqrt(vcov[5, 5] + 
                                  (1)^2 * vcov[13, 13] + 
                                  2 * 1 * vcov[5, 13]) # SE for marginal effect of predictor

national_case_me_ci_dem <- cbind(national_case_values * (1.96 * national_case_me_se_dem), 
                                 national_case_values * -(1.96 * national_case_me_se_dem))
national_case_me_ci_rep <- cbind(national_case_values * (1.96 * national_case_me_se_rep), 
                                 national_case_values * -(1.96 * national_case_me_se_rep))

national_case_preds <- national_case_pop_effect$fit 
national_case_upper <- national_case_pop_effect$fit + c(national_case_me_ci_dem[, 1], 
                                                        national_case_me_ci_rep[, 1])
national_case_lower <- national_case_pop_effect$fit + c(national_case_me_ci_dem[, 2], 
                                                        national_case_me_ci_rep[, 2])
national_case_final <- data.frame(Party = c(rep('Democrat', length(national_case_values)), 
                                            rep('Republican', length(national_case_values))), 
                                  X = as.numeric(national_case_values), 
                                  Y = as.numeric(national_case_preds), 
                                  U = as.numeric(national_case_upper), 
                                  L = as.numeric(national_case_lower),
                                  Y_linear = exp(as.numeric(national_case_preds)) - 1, 
                                  U_linear = exp(as.numeric(national_case_upper)) - 1, 
                                  L_linear = exp(as.numeric(national_case_lower)) - 1)

national_case_final$Party <- as.factor(national_case_final$Party)
national_case_final$Party <- factor(national_case_final$Party, 
                                    levels = rev(levels(national_case_final$Party)))

ggplot(national_case_final, 
       aes(x = X, 
           y = Y_linear, 
           color = Party)) + 
  geom_line() +
  theme_calc() +
  geom_ribbon(aes(ymin = L_linear, 
                  ymax = U_linear), 
              linetype = 2, 
              alpha = 0.0) +
  theme(text = element_text(size = 12.5), 
        plot.title = element_text(hjust = 0.5),
        plot.margin = grid::unit(c(5, 5, 5, 5), "mm"),
        legend.title = element_blank(),
        legend.position = 'bottom',
        legend.direction = 'horizontal') +
  scale_color_manual(values = c('darkred', 'dodgerblue4')) +
  xlab('\nWeekly Count of National New Cases (per 10k)') + 
  ylab('Count of COVID-19 Tweets\n') +
  geom_rug(data = df_uni, 
           aes(x = national_case_pop), 
           inherit.aes = FALSE, 
           color = 'gray60') +
  scale_shape_cleveland(overlap = FALSE)

ggsave(paste0(rep_path, 'FigS11c.png'),
       dpi = 600,
       width = 5,
       height = 3.5,
       units = 'in')


# national death -------------------------

national_death_values <- seq(round(min(df_uni$national_death_pop), 1),
                             round(max(df_uni$national_death_pop), 1),
                             0.1)
national_death_pop_effect <- Effect(focal.predictors = c('national_death_pop',
                                                         'party_new'), 
                                    mod = uni_all_inter,
                                    xlevels = list(national_death_pop = national_death_values),
                                    given.values = c(
                                      state_case_pop = median(df_uni$state_case_pop),
                                      state_death_pop = median(df_uni$state_death_pop),
                                      national_case_pop = median(df_uni$national_case_pop),
                                      state_covid_policy = median(df_uni$state_covid_policy),
                                      week_new = median(df_uni$week_new),
                                      week_new_2 = median(df_uni$week_new_2),
                                      week_new_3 = median(df_uni$week_new_3)
                                      )
                                    )

national_death_me_se_dem <- sqrt(vcov[6, 6]) # SE for marginal effect of predictor
national_death_me_se_rep <- sqrt(vcov[6, 6] + 
                                   (1)^2 * vcov[14, 14] + 
                                   2 * 1 * vcov[6, 14]) # SE for marginal effect of predictor

national_death_me_ci_dem <- cbind(national_death_values * (1.96 * national_death_me_se_dem), 
                                  national_death_values * -(1.96 * national_death_me_se_dem))
national_death_me_ci_rep <- cbind(national_death_values * (1.96 * national_death_me_se_rep), 
                                  national_death_values * -(1.96 * national_death_me_se_rep))

national_death_preds <- national_death_pop_effect$fit 
national_death_upper <- national_death_pop_effect$fit + c(national_death_me_ci_dem[, 1], 
                                                          national_death_me_ci_rep[, 1])
national_death_lower <- national_death_pop_effect$fit + c(national_death_me_ci_dem[, 2], 
                                                          national_death_me_ci_rep[, 2])
national_death_final <- data.frame(Party = c(rep('Democrat', length(national_death_values)), 
                                             rep('Republican', length(national_death_values))), 
                                   X = as.numeric(national_death_values), 
                                   Y = as.numeric(national_death_preds), 
                                   U = as.numeric(national_death_upper), 
                                   L = as.numeric(national_death_lower),
                                   Y_linear = exp(as.numeric(national_death_preds)) - 1, 
                                   U_linear = exp(as.numeric(national_death_upper)) - 1, 
                                   L_linear = exp(as.numeric(national_death_lower)) - 1)

national_death_final$Party <- as.factor(national_death_final$Party)
national_death_final$Party <- factor(national_death_final$Party, 
                                     levels = rev(levels(national_death_final$Party)))

ggplot(national_death_final,
       aes(x = X, 
           y = Y_linear, 
           color = Party)) + 
  theme_calc() +
  geom_line() +
  theme(text = element_text(size = 12.5), 
        plot.title = element_text(hjust = 0.5),
        plot.margin = grid::unit(c(5, 5, 5, 5), "mm"),
        legend.title = element_blank(),
        legend.position = 'bottom',
        legend.direction = 'horizontal') +
  scale_color_manual(values = c('darkred', 'dodgerblue4')) +
  geom_ribbon(aes(ymin = L_linear, ymax = U_linear), 
              linetype = 2, 
              alpha = 0.0) +
  xlab('\nWeekly Count of National New Deaths (per 10k)') + 
  ylab('Count of COVID-19 Tweets\n') +
  geom_rug(data = df_uni, 
           aes(x = national_death_pop), 
           inherit.aes = FALSE, 
           color = 'gray60') +
  scale_shape_cleveland(overlap = FALSE)

ggsave(paste0(rep_path, 'FigS11d.png'),
       dpi = 600,
       width = 5,
       height = 3.5,
       units = 'in')


# state policy -------------------------

state_policy_values <- seq(round(min(df_uni$state_covid_policy), 1),
                           round(max(df_uni$state_covid_policy), 1),
                           0.01)
state_policy_effect <- Effect(focal.predictors = c('state_covid_policy', 
                                                   'party_new'), 
                              mod = uni_all_inter,
                              xlevels = list(state_covid_policy = state_policy_values),
                              given.values = c(
                                state_case_pop = median(df_uni$state_case_pop),
                                state_death_pop = median(df_uni$state_death_pop),
                                national_case_pop = median(df_uni$national_case_pop),
                                national_death_pop = median(df_uni$national_death_pop),
                                week_new = median(df_uni$week_new),
                                week_new_2 = median(df_uni$week_new_2),
                                week_new_3 = median(df_uni$week_new_3)
                                )
                              )

state_policy_me_se_dem <- sqrt(vcov[7, 7]) # SE for marginal effect of predictor
state_policy_me_se_rep <- sqrt(vcov[7, 7] + 
                                 (1)^2 * vcov[15, 15] + 
                                 2 * 1 * vcov[7, 15]) # SE for marginal effect of predictor

state_policy_me_ci_dem <- cbind(state_policy_values * (1.96 * state_policy_me_se_dem), 
                                state_policy_values * -(1.96 * state_policy_me_se_dem))
state_policy_me_ci_rep <- cbind(state_policy_values * (1.96 * state_policy_me_se_rep), 
                                state_policy_values * -(1.96 * state_policy_me_se_rep))

state_policy_preds <- state_policy_effect$fit
state_policy_upper <- state_policy_effect$fit + c(state_policy_me_ci_dem[, 1], 
                                                  state_policy_me_ci_rep[, 1])
state_policy_lower <- state_policy_effect$fit + c(state_policy_me_ci_dem[, 2], 
                                                  state_policy_me_ci_rep[, 2])
state_policy_final <- data.frame(Party = c(rep('Democrat',length(state_policy_values)), 
                                           rep('Republican',length(state_policy_values))), 
                                 X = as.numeric(state_policy_values), 
                                 Y = as.numeric(state_policy_preds), 
                                 U = as.numeric(state_policy_upper), 
                                 L = as.numeric(state_policy_lower),
                                 Y_linear = exp(as.numeric(state_policy_preds)) - 1, 
                                 U_linear = exp(as.numeric(state_policy_upper)) - 1, 
                                 L_linear = exp(as.numeric(state_policy_lower)) - 1)

state_policy_final$Party <- as.factor(state_policy_final$Party)
state_policy_final$Party <- factor(state_policy_final$Party, 
                                   levels = rev(levels(state_policy_final$Party)))

ggplot(state_policy_final, 
       aes(x = X, 
           y = Y_linear, 
           color = Party)) + 
  theme_calc() +
  geom_line() +
  theme(text = element_text(size = 12.5), 
        plot.title = element_text(hjust = 0.5),
        plot.margin = grid::unit(c(5, 5, 5, 5), "mm"),
        legend.title = element_blank(),
        legend.position='bottom',
        legend.direction='horizontal') +
  scale_x_continuous(breaks = seq(0, 11, by = 1))+
  scale_y_continuous(limits = c(0, 1.5)) +
  scale_color_manual(values = c('darkred', 'dodgerblue4')) +
  geom_ribbon(aes(ymin = L_linear, ymax = U_linear), 
              linetype = 2, 
              alpha = 0.0) +
  xlab('\nWeekly Count of State New COVID-19 Policies') + 
  ylab('Count of COVID-19 Tweets\n') +
  geom_histogram(
    data = df_uni, 
    aes(x = state_covid_policy, y = ..density..), 
    binwidth = 1,
    inherit.aes = FALSE, 
    fill = 'gray', 
    alpha = 0.35) +
  scale_shape_cleveland(overlap = FALSE)

ggsave(paste0(rep_path, 'FigS11e.png'),
       dpi = 600,
       width = 5,
       height = 3.5,
       units = 'in')
